"""
function: adjust images size and ground truth size
author: lyk
create time: 2021.10.15
"""
import os
import shutil
import cv2
from tqdm import tqdm
import xml.etree.ElementTree as ET


def read_image_list(file_list):
    """
    input file_list: file.txt containing absolute path
    output: list containing absolute path
    """
    file_collection = []
    with open(file_list, 'r') as file_r:
        lines = file_r.readlines()
        for line in lines:
            value = line.strip()
            file_collection.append(value)
    return file_collection


def adjust_size(collection_files, src_dataset_name, dst_dataset_name):
    """
    input collection_files: list containing absolute path
    input src_dataset_name: src dataset name
    input dst_dataset_name: dst dataset name
    """
    for index_file in tqdm(collection_files):
        if not os.path.exists(index_file):
            print(index_file)
            continue
        # test_file = 'D:\\00_indemind_lyk\\tinyDataSet\\tinyDataSet\images\\train\\000000000030_coco.jpg'
        # test_file = 'D:\\00_indemind_lyk\\tinyDataSet\\tinyDataSet\images\\train\\000000000073_coco.jpg'
        # test_file = 'D:\\00_indemind_lyk\\tinyDataSet\\tinyDataSet\images\\train\\cam12021_05_29_16_52_37_1622307157225.jpg'
        image = cv2.imread(index_file)
        image_shape = image.shape
        index_min = image_shape.index(min([image_shape[0], image_shape[1]]))

        ratio = 320 / image_shape[index_min]
        resize_image = cv2.resize(image, (0, 0), fx=ratio, fy=ratio)
        dst_file = index_file.replace(src_dataset_name, dst_dataset_name)
        dst_folder, basename = os.path.split(dst_file)
        if not os.path.exists(dst_folder):
            os.makedirs(dst_folder)
        # print(dst_file)
        cv2.imwrite(dst_file, resize_image)


def adjust_xml(collection_files, src_dataset_name, anno_dataset_name, dst_dataset_name):
    """
    input collection_files: list containing absolute path
    input src_dataset_name: src dataset name
    input dst_dataset_name: dst dataset name
    """
    for index_file in tqdm(collection_files):
        anno_index_file = index_file.replace(src_dataset_name, anno_dataset_name)[:-4] + '.xml'

        if not os.path.exists(anno_index_file):
            print(index_file)
            continue
        tree = ET.parse(anno_index_file)
        root = tree.getroot()
        image_size_xml = root.find('size')

        image_width = image_size_xml.find('width')
        image_height = image_size_xml.find('height')
        image_w = int(image_width.text)
        image_h = int(image_height.text)
        scale = 320 / min(image_w, image_h)
        img_w = int(image_w * scale)
        img_h = int(image_h * scale)
        image_width.text = str(img_w)
        image_height.text = str(img_h)
        objects = root.findall('object')
        for idx, object in enumerate(objects):
            bndbox = object.find('bndbox')
            xmin = bndbox.find('xmin')
            xmax = bndbox.find('xmax')
            ymin = bndbox.find('ymin')
            ymax = bndbox.find('ymax')
            if 0 < int(xmin.text) < 1 and 0 < int(ymin.text) < 1 and 0 < int(ymin.text) < 1 and 0 < int(ymax.text) < 1:
                continue
            xmin.text = str(int(int(xmin.text) * scale))
            xmax.text = str(int(int(xmax.text) * scale))
            ymin.text = str(int(int(ymin.text) * scale))
            ymax.text = str(int(int(ymax.text) * scale))

        dst_index_file = anno_index_file.replace(anno_dataset_name, dst_dataset_name)
        dst_folder, basename = os.path.split(dst_index_file)
        if not os.path.exists(dst_folder):
            os.makedirs(dst_folder)
        tree.write(dst_index_file, method='xml', encoding='utf-8')


if __name__ == '__main__':
    index_train_file_list = 'D:/00_indemind_lyk/tinyDataSet/tinyDataSet/dir_train.txt'
    index_val_file_list = 'D:/00_indemind_lyk/tinyDataSet/tinyDataSet/dir_val.txt'
    index_train_file_collection = read_image_list(index_train_file_list)
    index_val_file_collection = read_image_list(index_val_file_list)
    index_file_collection = []
    index_file_collection.extend(index_train_file_collection)
    index_file_collection.extend(index_val_file_collection)
    # adjust_size(index_file_collection, 'images', 'images_res')
    adjust_xml(index_file_collection, 'images', 'xml_labels', 'xml_res')